import time
import datetime
from bpfcc import BPF, USDT
import json
import atexit
import platform
import subprocess
import sys
import os

# Define the BPF program
bpf_text = """
#include <uapi/linux/ptrace.h>
#include <linux/sched.h>

#define MAX_ENTRIES 1024

enum lock_type_e {
    mutex_lock = 1,
    mutex_unlock,
    mutex_wrlock,
    mutex_rdlock,
    mutex_timedlock,
    mutex_rwlock_unlock,
    sem_wait,
    spin_lock
};

struct event {
    u64 start_ts;
    u64 end_ts;
    u32 pid;
    u32 tid;
    u64 lock;
    enum lock_type_e lock_type;
    char comm[TASK_COMM_LEN];
};

struct trace_lock {
    u64 tid;
    u64 start_ts;
    u64 lock_p;
};

BPF_HASH(start_times, u64, struct trace_lock);
BPF_PERF_OUTPUT(events);

int trace_lock_entry(struct pt_regs *ctx)
{
    struct trace_lock lock = {};

    u64 ts = bpf_ktime_get_ns();
    u64 tid = (u64)bpf_get_current_pid_tgid();

    lock.tid = tid;
    lock.start_ts = ts;
    lock.lock_p = (u64)PT_REGS_PARM1(ctx);

    // Store the timestamp in the hash map
    start_times.update(&tid, &lock);

    return 0;
}

static int trace_lock_exit_common(enum lock_type_e lock_type, struct pt_regs *ctx)
{
    struct event e = {};
    struct lock *lock = 0;
    u64 ts = 0;
    u64 tid = (u64)bpf_get_current_pid_tgid();

    // Get current timestamp in nanoseconds
    ts = bpf_ktime_get_ns();

    // Retrieve the start time from the hash map
    lock = start_times.lookup(&tid);
    if (!lock)
        return 0;   // Missed tracing context

    // Populate the event structure
    e.start_ts = lock->start_ts;
    e.end_ts = ts;
    e.pid = tid >> 32;
    e.tid = (u32)tid;
    e.lock_type = lock_type;
    e.lock = lock->lock_p;

    // Remove the entry from the hash map
    start_times.delete(&tid);

    // Output the event to user space
    DELAY_FILTER
    {
        bpf_get_current_comm(&e.comm, sizeof(e.comm));
        events.perf_submit(ctx, &e, sizeof(e));
    }
    return 0;
} 

int trace_mutex_lock_exit(struct pt_regs *ctx)
{
    return trace_lock_exit_common(mutex_lock, ctx);
}

int trace_mutex_unlock_exit(struct pt_regs *ctx)
{
    return trace_lock_exit_common(mutex_unlock, ctx);
}

int trace_sem_wait_exit(struct pt_regs *ctx)
{
    return trace_lock_exit_common(sem_wait, ctx);
}

int trace_spin_lock_exit(struct pt_regs *ctx)
{
    return trace_lock_exit_common(spin_lock, ctx);
}

int trace_mutex_wrlock_exit(struct pt_regs *ctx)
{
    return trace_lock_exit_common(mutex_wrlock, ctx);
}

int trace_mutex_rdlock_exit(struct pt_regs *ctx)
{
    return trace_lock_exit_common(mutex_rdlock, ctx);
}

int trace_mutex_timedlock_exit(struct pt_regs *ctx)
{
    return trace_lock_exit_common(mutex_timedlock, ctx);
}

int trace_rwlock_unlock_exit(struct pt_regs *ctx)
{
    return trace_lock_exit_common(mutex_rwlock_unlock, ctx);
}

"""
def _convert_ktime_to_unix(bpf_ktime_ns):
    # 获取当前时间戳（CLOCK_REALTIME）和系统运行时间（uptime）
    current_realtime = time.time()
    with open("/proc/uptime", "r") as f:
        uptime_seconds = float(f.read().split()[0])
    
    # 计算系统启动时间（CLOCK_REALTIME at boot）
    boot_time = current_realtime - uptime_seconds
    
    # 将 eBPF 时间戳转换为 Unix 时间戳（纳秒级）
    unix_ns = boot_time * 1e9 + bpf_ktime_ns
    return unix_ns

def _get_match_runtime_path(so_path, substr):
    for s in so_path:
        if substr in s:
            return s
    return None

def _get_host_runtime_path(directory, substr):
    """
    在指定目录及其子目录中递归查找 libruntime.so 文件。

    :param directory: 要搜索的目录路径
    :return: 找到的第一个 libruntime.so 文件的绝对路径，如果没有找到则返回 None
    """
    so_path = []
    for root, dirs, files in os.walk(directory):
        if 'libruntime.so' in files:
            so_path.append(os.path.join(root, 'libruntime.so'))
    
    return _get_match_runtime_path(so_path, substr)

def _get_container_runtime_path(container_id, substr):
    try:
        # 获取容器的根文件系统路径
        inspect_command = f"docker inspect --format '{{{{.GraphDriver.Data.MergedDir}}}}' {container_id}"
        container_root = subprocess.check_output(inspect_command, shell=True).decode().strip()

        # 查找容器内的 libruntime.so文件
        find_command = f"find {container_root} -name libruntime.so"
        so_path = subprocess.check_output(find_command, shell=True).decode().strip()

        if so_path:
            return _get_match_runtime_path(so_path.split("\n"), substr)
        else:
            return "libruntime.so not found in the container."

    except subprocess.CalledProcessError as e:
        return f"Error: {e.output.decode().strip()}"

class FunctionName(FuncType):
    MUTEX_LOCK = 1
    MUTEX_UNLOCK = 2
    MUTEX_WRLOCK = 3
    MUTEX_RDLOCK = 4
    MUTEX_TIMEDLOCK = 5
    MUTEX_RWLOCK_UNLOCK = 6
    SEM_WAIT = 7
    SPIN_LOCK = 8

def get_function_name(func_type: int) -> str:
    try:
        func_enum = FunctionName(func_type)
        # 根据枚举成员返回预定义的字符串
        if func_enum == FunctionName.MUTEX_LOCK:
            return "mutex_lock"
        elif func_enum == FunctionName.MUTEX_UNLOCK:
            return "mutex_unlock"
        elif func_enum == FunctionName.MUTEX_WRLOCK:
            return "rwlock_wrlock"
        elif func_enum == FunctionName.MUTEX_RDLOCK:
            return "rwlock_rdlock"
        elif func_enum == FunctionName.MUTEX_TIMEDLOCK:
            return "mutex_timedlock"
        elif func_enum == FunctionName.MUTEX_RWLOCK_UNLOCK:
            return "rwlock_unlock"
        elif func_enum == FunctionName.SEM_WAIT:
            return "sem_wait"
        elif func_enum == FunctionName.SPIN_LOCK:
            return "spin_lock"
    except ValueError:
        # 如果 status_code 不是 Status 枚举的有效值
        pass

        return "Unknown"

# Initialize BPF
b = BPF(text=bpf_text)

path_or_container = ""

if len(sys.argv) != 4:
   raise SystemExit(0)

path_or_container = sys.argv[1]
match_str = sys.argv[2]
delay = (u64)sys.argv[3]
delay = delay * 1000 * 1000

bpf_text = bpf_text.replace('DELAY_FILTER',
        'if ((e.end_ts - e.start_ts) > %lld)' % delay)

if path_or_container[0] == '.' or path_or_container[0] == '/':
   lib_path = _get_host_runtime_path(path_or_container, match_str)
else:
   lib_path = _get_container_runtime_path(path_or_container, match_str)

print("Hook so: ", lib_path)

# Attach uprobes to pthread_mutex_lock/pthread_mutex_unlock ---
b.attach_uprobe(name=lib_path, sym="pthread_mutex_lock", fn_name="trace_lock_entry")
b.attach_uretprobe(name=lib_path, sym="pthread_mutex_lock", fn_name="trace_mutex_lock_exit")

b.attach_uprobe(name=lib_path, sym="pthread_mutex_unlock", fn_name="trace_lock_entry")
b.attach_uretprobe(name=lib_path, sym="pthread_mutex_unlock", fn_name="trace_mutex_unlock_exit")

b.attach_uprobe(name=lib_path, sym="sem_wait", fn_name="trace_lock_entry")
b.attach_uretprobe(name=lib_path, sym="sem_wait", fn_name="trace_sem_wait_exit")

b.attach_uprobe(name=lib_path, sym="pthread_spin_lock", fn_name="trace_lock_entry")
b.attach_uretprobe(name=lib_path, sym="pthread_spin_lock", fn_name="trace_spin_lock_exit")

b.attach_uprobe(name=lib_path, sym="pthread_rwlock_wrlock", fn_name="trace_lock_entry")
b.attach_uretprobe(name=lib_path, sym="pthread_rwlock_wrlock", fn_name="trace_mutex_wrlock_exit")

b.attach_uprobe(name=lib_path, sym="pthread_rwlock_rdlock", fn_name="trace_lock_entry")
b.attach_uretprobe(name=lib_path, sym="pthread_rwlock_rdlock", fn_name="trace_mutex_rdlock_exit")

b.attach_uprobe(name=lib_path, sym="pthread_mutex_timedlock", fn_name="trace_lock_entry")
b.attach_uretprobe(name=lib_path, sym="pthread_mutex_timedlock", fn_name="trace_mutex_timedlock_exit")

b.attach_uprobe(name=lib_path, sym="pthread_rwlock_unlock", fn_name="trace_lock_entry")
b.attach_uretprobe(name=lib_path, sym="pthread_rwlock_unlock", fn_name="trace_rwlock_unlock_exit")

# List to store events
events = []

def print_event(cpu, data, size):
    event = b["events"].event(data)
    decoded_func_name = get_function_name(event.lock_type)
    chrome_trace_event = {
        "name": decoded_func_name,
        "cat": "lock",
        "ph": "X",
        "pid": event.pid,
        "tid": event.tid,
        "ts": _convert_ktime_to_unix(event.start_ts) / 1000.0, # Convert to microseconds
        "dur": (event.end_ts - event.start_ts) / 1000.0,       # Duration in microseconds
        "args": {"comm": event.comm.decode('utf-8'), "lock": event.lock}
    }
    events.append(chrome_trace_event)

# Register the callback function
b["events"].open_perf_buffer(print_event)

def dump_events():
    with open("log.json", "w") as f:
        json.dump(events, f, indent=2)
    print("Events dumped to log.json")

atexit.register(dump_events)

print("Tracing lock... Hit Ctrl-C to end.")

# Loop to process events
while True:
    try:
        b.perf_buffer_poll()
    except KeyboardInterrupt:
        break

# useage1: python3 mutex_trace.py  <container_id>  "libruntime.so" <max_delay_ms>
# useage2: python3 mutex_trace.py  /  "libruntime.so" <max_delay_ms>